# -*- coding:utf-8 -*-
import csv
import os

from options import HiDDenConfiguration
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torchvision
import utils
from model.discriminator import Discriminator
import numpy as np
from model.encoder_decoder import EncoderDecoder
def denormalize(tensor):
    """ 反归一化：从 [-1, 1] 转换回 [0,1] """
    return torch.clamp((tensor + 1) / 2, 0, 1)

def test(train_loader, test_loader, messages_length, num_epochs, model_dir, log_file, save_image_dir, en, de, adv, normalize):

    hidden_config = HiDDenConfiguration(H=128, W=128,
                                        message_length=messages_length,
                                        encoder_blocks=4, encoder_channels=64,
                                        decoder_blocks=7, decoder_channels=64,
                                        use_discriminator=True,
                                        use_vgg=False,
                                        discriminator_blocks=3, discriminator_channels=64,
                                        decoder_loss=1,
                                        encoder_loss=0.7,
                                        adversarial_loss=1e-3
                                        )
    mse_loss = nn.MSELoss().cuda()
    bce_with_logits_loss = nn.BCEWithLogitsLoss()
    discriminator = Discriminator(hidden_config).cuda()

    encoder_decoder = EncoderDecoder(hidden_config)
    encoder_decoder.load_state_dict(torch.load("/home/zzc/simulator_noisy/runs-works/run_base_2025-04-22_21-26-54/model/300_model.pth"))
    encoder_decoder.cuda()
    
    encoder_decoder.noiser.load_state_dict(torch.load('/home/zzc/simulator_noisy/run-UnetPlusPlusL6_mse_dct64_2025-04-20_20-43-40/model/best_model.pth'))
    encoder_decoder.noiser.eval()
    cover_data_label = 1

    test_running_loss = 0.0
    dis_running_loss = 0.0

    with open(log_file, "w") as f:
        writer = csv.writer(f)
        writer.writerow(["testLoss", "advLoss", "bit_arr", "psnr", "ssim"])
        test_psnr = 0.0
        test_ssim = 0.0
        test_bit_arr = 0.0
        for images, _ in test_loader:
            images = images.cuda()
            target_label_adv = torch.full((images.size(0), 1), cover_data_label).cuda()

            messages = torch.Tensor(np.random.choice([0, 1], size=[images.size(0), 1, messages_length, messages_length])).cuda()
            encoder_image, noise_image, decoder_message = encoder_decoder(images, messages)
            # 测试阶段
            encoder_decoder.eval()
            discriminator.eval()
            with torch.no_grad():
                dis_loss = bce_with_logits_loss(discriminator(encoder_image), target_label_adv.float())
                encoder_loss = mse_loss(images, encoder_image)
                decoder_loss = mse_loss(messages, decoder_message)
                loss = en * encoder_loss + de * decoder_loss + adv * dis_loss

                psnr, ssim = utils.val_psnr_ssim(images, encoder_image)
                test_psnr += psnr
                test_ssim += ssim
                test_running_loss += loss.item()
                dis_running_loss += dis_loss.item()


                save_images(images.cpu()[:8, :, :, :],
                            encoder_image[:8, :, :, :].cpu(), noise_image[:8, :, :, :].cpu(), 0, save_image_dir, normalize)
                decoded_rounded = decoder_message.detach().cpu().numpy().round().clip(0, 1)
                bitwise_err = np.sum(np.abs(decoded_rounded - messages.detach().cpu().numpy())) / (messages.shape[0] * messages.shape[1]* messages.shape[2] * messages.shape[3])
                test_bit_arr += bitwise_err

        test_epoch_loss = test_running_loss / len(test_loader)
        dis_epoch_loss = dis_running_loss / len(test_loader)
        psnr = test_psnr / len(test_loader)
        ssim = test_ssim / len(test_loader)
        bit_arr = test_bit_arr / len(test_loader)

        # 格式化损失为保留4位小数的字符串
        test_loss_str = "{:.4f}".format(test_epoch_loss)
        dis_loss_str = "{:.4f}".format(dis_epoch_loss)
        psnr = "{:.2f}".format(psnr)
        ssim = "{:.2f}".format(ssim)

        writer.writerow([test_loss_str, dis_loss_str, bit_arr, psnr, ssim])

        print(f"Test Loss: {test_epoch_loss:.4f}, Adv Loss：{dis_epoch_loss}, Bit arr：{bit_arr}, Psnr：{psnr}, Ssim：{ssim}")


def save_images(cover_images, encoder_images, noise_image, epoch, folder, normalize):
    # 将数据转移到 CPU
    images = cover_images.cpu()
    encoder_images = encoder_images.cpu()
    noise_image = noise_image.cpu()

    # 反归一化
    if normalize:
        images = denormalize(images)
        encoder_images = denormalize(encoder_images)
        noise_image = denormalize(noise_image)

    # 计算误差放大
    revise_image = torch.clamp(torch.abs(images - encoder_images) * 5, 0, 1)
    revise_image_5 = torch.clamp(torch.abs(images - noise_image) * 5, 0, 1)

    # 拼接图像
    stacked_images = torch.cat([images, encoder_images, noise_image, revise_image, revise_image_5], dim=0)

    # 保存图片
    filename = os.path.join(folder, 'epoch-{}.png'.format(epoch + 1))
    torchvision.utils.save_image(stacked_images, filename, normalize=False)